#!/bin/bash

type function > /dev/null 2>&1
[ 0 -ne $? ] && echo 'Run with bash.' && exit 3
[ 1 -lt $# ] && echo "Usage: $0 [seconds]" && exit 1
[[ ! "$2" =~ ^[0-9]*$ ]] && "Usage: $0 [seconds]" && exit 2

CONF='/tmp/ipflow.conf'

function deal_forward {
    for ip in $(cat $CONF); do
        iptables -$1 FORWARD -s $ip -j COLBEN_UPLOAD
        iptables -$1 FORWARD -d $ip -j COLBEN_DOWNLOAD
    done
}

function insert_forward {
    cat /proc/net/arp | grep '^10' | grep -v '00:00:00:00:00:00' | awk '{print $1}' > $CONF
    deal_forward I
}

function clean_forward {
    [ ! -f "$CONF" ] && return 0
    deal_forward D
    rm -f $CONF
}

function show_result {
    clear
    echo 'Download:'
    iptables -nvxL FORWARD | grep COLBEN_DOWNLOAD | awk "{printf(\"%-4dKB/s:%-10s \n\",\$2/1024/$1,\$9)}" | sort -t: -r
    echo -e '\033[s\033[1;25HUpload:'
    line_num=1
    iptables -nvxL FORWARD | grep COLBEN_UPLOAD | awk "{printf(\"%-4dKB/s:%-10s\n\",\$2/1024/$1,\$8)}" | sort -t: -r | while read line; do
        let "line_num=1+$line_num"
        echo -e "\033[$line_num;25H$line"
    done
    echo -e '\033[u'
    sleep $1
}

function quit {
    clean_forward
    iptables -X COLBEN_UPLOAD
    iptables -X COLBEN_DOWNLOAD
    exit 0
}

iptables -N COLBEN_UPLOAD
iptables -N COLBEN_DOWNLOAD

trap "quit" 2 3 15

while :; do
    insert_forward
    iptables -Z FORWARD
    show_result ${1:-2}
    clean_forward
    sleep 1
done

